import cv2
import requests
import time
import threading
import socketio
import logging

from lerobot.common.robot_devices.control_configs import ControlPipelineConfig
from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot
import draccus
import torch


cli_args = ['--robot.type=so101', '--robot.cameras={}', '--control.type=teleoperate']

cfg = draccus.parse(config_class=ControlPipelineConfig, config_path=None, args=cli_args)
robot = ManipulatorRobot(cfg.robot)


if not robot.is_connected:
    robot.connect()


# 启用详细日志
logging.basicConfig(level=logging.INFO)

# SERVER_URL = "http://127.0.0.1:5000"
SERVER_URL = "http://47.109.17.68:8868"

# 创建SocketIO客户端
sio = socketio.Client(logger=True, engineio_logger=True)

def send_frame(frame):
    """发送帧到服务器"""
    _, buffer = cv2.imencode('.jpg', frame)
    frame_bytes = buffer.tobytes()
    try:
        response = requests.post(SERVER_URL + '/receive_frame', 
                               data=frame_bytes,
                               headers={'Content-Type': 'image/jpeg'},
                               timeout=2)
        if response.status_code != 200:
            print(f"Frame sending failed with status: {response.status_code}")
    except requests.exceptions.RequestException as e:
        print(f"Error sending frame: {e}")

def capture_and_send():
    """捕获摄像头视频并发送"""
    cap = cv2.VideoCapture(10)
    if not cap.isOpened():
        print("Error: Could not open camera.")
        return
    
    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                print("Error: Could not read frame.")
                break
            
            # 调整帧大小以减少带宽
            frame = cv2.resize(frame, (640, 480))
            send_frame(frame)
            time.sleep(0.033)  # 约30fps
    except Exception as e:
        print(f"Camera error: {e}")
    finally:
        cap.release()
        print("Camera released")

# ========== 连接到 /video 命名空间 ==========
@sio.event
def connect():
    print('Connected to server on default namespace')

@sio.event
def disconnect():
    print('Disconnected from server')

@sio.on('numpy_data_to_video', namespace='/video')
def handle_numpy_data(data):
    """接收来自服务端的numpy数据"""
    print(f"Video client received numpy data: {data}")
    action = torch.tensor(data)
    robot.send_action(action)

if __name__ == '__main__':
    try:
        # 连接到服务端
        sio.connect(SERVER_URL, namespaces=['/video'])
        print(f"Connected to {SERVER_URL} on namespaces: /video")
        
        # 启动视频捕获和发送线程
        video_thread = threading.Thread(target=capture_and_send)
        video_thread.daemon = True
        video_thread.start()
        
        # 保持主线程运行
        while True:
            time.sleep(1)
    except socketio.exceptions.ConnectionError as e:
        print(f"Connection failed: {e}")
    except KeyboardInterrupt:
        sio.disconnect()
        print("Disconnected")